"""
LangGraph 深度集成模块
将增强记忆功能深度集成到 Supervisor 工作流中
"""

from typing import Dict, Any, List, Optional, TypedDict, Annotated
from datetime import datetime
import logging

from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
import operator

from .enhanced_memory_class import EnhancedMemoryClass
from .query_rewriter import QueryRewriter

logger = logging.getLogger(__name__)

class EnhancedAgentState(TypedDict):
    """增强的智能体状态"""
    messages: Annotated[List[BaseMessage], operator.add]
    next: str
    user_id: str
    session_id: str
    context: Dict[str, Any]
    emotion: Dict[str, Any]
    
    # 增强记忆相关字段
    original_query: str
    rewritten_query: str
    memory_context: Dict[str, Any]
    user_profile: Dict[str, Any]
    conversation_history: List[Dict[str, str]]
    relevant_memories: List[Dict[str, Any]]
    rewrite_info: Optional[Dict[str, Any]]

class MemoryEnhancedNode:
    """记忆增强节点基类"""
    
    def __init__(self, enhanced_memory: EnhancedMemoryClass):
        self.enhanced_memory = enhanced_memory
    
    async def process_with_memory(
        self, 
        state: EnhancedAgentState, 
        node_name: str
    ) -> Dict[str, Any]:
        """使用记忆增强处理"""
        try:
            # 确保记忆系统已初始化
            if not self.enhanced_memory._initialized:
                await self.enhanced_memory.initialize()
            
            # 获取增强上下文
            if state.get("messages"):
                current_message = state["messages"][-1]
                if isinstance(current_message, HumanMessage):
                    current_query = current_message.content
                    
                    # 获取增强上下文
                    memory_context = await self.enhanced_memory.get_enhanced_context(
                        current_query,
                        state.get("user_id", "default"),
                        state.get("session_id", "default")
                    )
                    
                    # 更新状态
                    state.update({
                        "original_query": memory_context["original_query"],
                        "rewritten_query": memory_context["rewritten_query"],
                        "memory_context": memory_context,
                        "user_profile": memory_context["user_profile"],
                        "conversation_history": memory_context["conversation_history"],
                        "relevant_memories": memory_context["relevant_memories"],
                        "rewrite_info": memory_context["rewrite_info"]
                    })
            
            return state
            
        except Exception as e:
            logger.error(f"Failed to process with memory in {node_name}: {e}")
            return state

class QueryRewriteNode(MemoryEnhancedNode):
    """查询重写节点"""
    
    def __init__(self, enhanced_memory: EnhancedMemoryClass):
        super().__init__(enhanced_memory)
        self.query_rewriter = QueryRewriter()
    
    async def __call__(self, state: EnhancedAgentState) -> EnhancedAgentState:
        """执行查询重写"""
        try:
            logger.info("Executing query rewrite node")
            
            # 处理记忆增强
            state = await self.process_with_memory(state, "query_rewrite")
            
            # 如果有重写信息，添加到消息中
            if state.get("rewrite_info") and state["rewrite_info"]["confidence"] > 0.7:
                rewrite_message = AIMessage(
                    content=f"[查询重写] 原查询: {state['original_query']} -> 重写后: {state['rewritten_query']}",
                    name="query_rewriter"
                )
                state["messages"].append(rewrite_message)
                
                logger.info(f"Query rewritten: {state['original_query']} -> {state['rewritten_query']}")
            
            return state
            
        except Exception as e:
            logger.error(f"Query rewrite node failed: {e}")
            return state

class MemoryContextNode(MemoryEnhancedNode):
    """记忆上下文节点"""
    
    async def __call__(self, state: EnhancedAgentState) -> EnhancedAgentState:
        """注入记忆上下文"""
        try:
            logger.info("Executing memory context node")
            
            # 处理记忆增强
            state = await self.process_with_memory(state, "memory_context")
            
            # 构建上下文消息
            context_parts = []
            
            # 用户画像功能已移除 - 简化记忆系统
            # if state.get("user_profile"):
            #     profile = state["user_profile"]
            #     if profile.get("preferences"):
            #         context_parts.append(f"用户偏好: {profile['preferences']}")
            #     if profile.get("interests"):
            #         context_parts.append(f"用户兴趣: {', '.join(profile['interests'])}")
            
            # 添加相关历史记忆
            if state.get("relevant_memories"):
                memory_summaries = [m["content"][:100] + "..." for m in state["relevant_memories"][:2]]
                if memory_summaries:
                    context_parts.append(f"相关历史: {'; '.join(memory_summaries)}")
            
            # 如果有上下文信息，添加到消息中
            if context_parts:
                context_message = AIMessage(
                    content=f"[上下文信息] {' | '.join(context_parts)}",
                    name="memory_context"
                )
                state["messages"].append(context_message)
                
                logger.info(f"Added memory context: {len(context_parts)} items")
            
            return state
            
        except Exception as e:
            logger.error(f"Memory context node failed: {e}")
            return state

class EnhancedSupervisorNode(MemoryEnhancedNode):
    """增强的 Supervisor 节点"""
    
    def __init__(self, enhanced_memory: EnhancedMemoryClass, original_supervisor):
        super().__init__(enhanced_memory)
        self.original_supervisor = original_supervisor
    
    async def __call__(self, state: EnhancedAgentState) -> EnhancedAgentState:
        """执行增强的 Supervisor 决策"""
        try:
            logger.info("Executing enhanced supervisor node")
            
            # 处理记忆增强
            state = await self.process_with_memory(state, "enhanced_supervisor")
            
            # 构建增强的提示上下文
            enhanced_context = self._build_enhanced_context(state)
            
            # 如果有增强上下文，添加到消息中
            if enhanced_context:
                context_message = AIMessage(
                    content=f"[Supervisor上下文] {enhanced_context}",
                    name="enhanced_supervisor"
                )
                # 在最后一条用户消息前插入上下文
                if state["messages"] and isinstance(state["messages"][-1], HumanMessage):
                    state["messages"].insert(-1, context_message)
                else:
                    state["messages"].append(context_message)
            
            # 调用原始 Supervisor
            result = await self.original_supervisor(state)
            
            return result
            
        except Exception as e:
            logger.error(f"Enhanced supervisor node failed: {e}")
            # 降级到原始 Supervisor
            return await self.original_supervisor(state)
    
    def _build_enhanced_context(self, state: EnhancedAgentState) -> str:
        """构建增强上下文"""
        context_parts = []
        
        # 查询重写信息
        if state.get("rewrite_info") and state["rewrite_info"]["confidence"] > 0.5:
            rewrite_info = state["rewrite_info"]
            context_parts.append(
                f"查询已重写(置信度:{rewrite_info['confidence']:.2f}): "
                f"{state.get('original_query', '')} -> {state.get('rewritten_query', '')}"
            )
        
        # 用户画像功能已移除 - 简化记忆系统
        # if state.get("user_profile"):
        #     profile = state["user_profile"]
        #     profile_summary = []
        #     if profile.get("preferences"):
        #         profile_summary.append(f"偏好:{list(profile['preferences'].keys())[:3]}")
        #     if profile.get("interests"):
        #         profile_summary.append(f"兴趣:{profile['interests'][:3]}")
        #
        #     if profile_summary:
        #         context_parts.append(f"用户画像: {', '.join(profile_summary)}")
        
        # 相关记忆
        if state.get("relevant_memories"):
            memory_count = len(state["relevant_memories"])
            context_parts.append(f"找到{memory_count}条相关历史记忆")
        
        return " | ".join(context_parts)

class Text2SQLEnhancedNode(MemoryEnhancedNode):
    """增强的 Text2SQL 节点"""
    
    def __init__(self, enhanced_memory: EnhancedMemoryClass, original_text2sql_agent):
        super().__init__(enhanced_memory)
        self.original_agent = original_text2sql_agent
    
    async def __call__(self, state: EnhancedAgentState) -> EnhancedAgentState:
        """执行增强的 Text2SQL 处理"""
        try:
            logger.info("Executing enhanced Text2SQL node")
            
            # 处理记忆增强
            state = await self.process_with_memory(state, "text2sql_enhanced")
            
            # 使用重写后的查询
            if state.get("rewritten_query") and state.get("rewrite_info"):
                rewrite_info = state["rewrite_info"]
                if rewrite_info["confidence"] > 0.7:
                    # 替换最后一条用户消息为重写后的查询
                    if state["messages"] and isinstance(state["messages"][-1], HumanMessage):
                        state["messages"][-1] = HumanMessage(
                            content=state["rewritten_query"]
                        )
                        logger.info(f"Using rewritten query for Text2SQL: {state['rewritten_query']}")
            
            # 调用原始 Text2SQL 智能体
            result = await self.original_agent(state)
            
            # 存储查询结果到记忆
            await self._store_query_result(state, result)
            
            return result
            
        except Exception as e:
            logger.error(f"Enhanced Text2SQL node failed: {e}")
            # 降级到原始智能体
            return await self.original_agent(state)
    
    async def _store_query_result(self, state: EnhancedAgentState, result: EnhancedAgentState):
        """存储查询结果到记忆"""
        try:
            if result.get("messages"):
                last_message = result["messages"][-1]
                if isinstance(last_message, AIMessage):
                    # 存储查询和结果
                    await self.enhanced_memory.store_conversation_turn(
                        user_message=state.get("rewritten_query", state.get("original_query", "")),
                        ai_response=last_message.content,
                        user_id=state.get("user_id", "default"),
                        session_id=state.get("session_id", "default"),
                        metadata={
                            "node_type": "text2sql",
                            "query_type": "database_query",
                            "rewrite_confidence": state.get("rewrite_info", {}).get("confidence", 0.0)
                        }
                    )
                    
                    # 触发摘要检查
                    await self.enhanced_memory.trigger_summary_if_needed(
                        state.get("user_id", "default"),
                        state.get("session_id", "default")
                    )
                    
        except Exception as e:
            logger.error(f"Failed to store query result: {e}")

class EnhancedWorkflowBuilder:
    """增强工作流构建器"""
    
    def __init__(self, enhanced_memory: EnhancedMemoryClass):
        self.enhanced_memory = enhanced_memory
    
    def create_enhanced_workflow(
        self, 
        original_supervisor, 
        original_agents: Dict[str, Any]
    ) -> StateGraph:
        """创建增强的工作流"""
        try:
            # 创建状态图
            workflow = StateGraph(EnhancedAgentState)
            
            # 创建增强节点
            query_rewrite_node = QueryRewriteNode(self.enhanced_memory)
            memory_context_node = MemoryContextNode(self.enhanced_memory)
            enhanced_supervisor = EnhancedSupervisorNode(self.enhanced_memory, original_supervisor)
            
            # 添加预处理节点
            workflow.add_node("query_rewrite", query_rewrite_node)
            workflow.add_node("memory_context", memory_context_node)
            workflow.add_node("enhanced_supervisor", enhanced_supervisor)
            
            # 添加增强的智能体节点
            enhanced_agents = {}
            for agent_name, agent in original_agents.items():
                if agent_name == "text2sql_agent":
                    enhanced_agent = Text2SQLEnhancedNode(self.enhanced_memory, agent)
                    enhanced_agents[agent_name] = enhanced_agent
                    workflow.add_node(agent_name, enhanced_agent)
                else:
                    # 其他智能体暂时使用原始版本
                    enhanced_agents[agent_name] = agent
                    workflow.add_node(agent_name, agent)
            
            # 设置工作流路径
            workflow.set_entry_point("query_rewrite")
            
            # 预处理流程
            workflow.add_edge("query_rewrite", "memory_context")
            workflow.add_edge("memory_context", "enhanced_supervisor")
            
            # Supervisor 路由
            workflow.add_conditional_edges(
                "enhanced_supervisor",
                lambda x: x["next"],
                {
                    "text2sql_agent": "text2sql_agent",
                    "FINISH": END
                }
            )
            
            # 智能体完成后的路由
            for agent_name in enhanced_agents.keys():
                workflow.add_edge(agent_name, "enhanced_supervisor")
            
            logger.info("Enhanced workflow created successfully")
            return workflow
            
        except Exception as e:
            logger.error(f"Failed to create enhanced workflow: {e}")
            raise
    
    def compile_enhanced_workflow(
        self, 
        workflow: StateGraph, 
        checkpointer: MemorySaver = None
    ):
        """编译增强工作流"""
        try:
            if checkpointer is None:
                checkpointer = MemorySaver()
            
            compiled_workflow = workflow.compile(checkpointer=checkpointer)
            logger.info("Enhanced workflow compiled successfully")
            return compiled_workflow
            
        except Exception as e:
            logger.error(f"Failed to compile enhanced workflow: {e}")
            raise
